from core.cuda import cuda_module
from core.optimizer import Optimizer


# 随机梯度下降法
class AdaGrad(Optimizer):
    def __init__(self, lr=0.001, eps=1e-8):
        """
        AdaGrad 优化器，用于自适应调整学习率。

        Args:
            lr (float): 学习率，控制参数更新的步长。
            eps (float): 为了数值稳定性而添加的小常数，一般取 1e-8。
        """
        super().__init__()
        self.lr = lr
        self.eps = eps
        self.hs = {}  # 存储梯度平方的累积和

    def update_one(self, param):
        """
        更新单个参数。

        Args:
            param (Parameter): 待更新的参数。

        Returns:
            None
        """
        xp = cuda_module

        h_key = id(param)
        if h_key not in self.hs:
            self.hs[h_key] = xp.zeros_like(param.data)

        lr = self.lr
        eps = self.eps
        grad = param.grad.data
        h = self.hs[h_key]

        h += grad * grad
        param.data -= lr * grad / (xp.sqrt(h) + eps)

